#pragma once

#include "UDP Test.h"
#include "../UDP Lib/RUdpSocket.h"
#include "../UDP Lib/ClientList.h"
#include "../UDP Lib/RUdpPacket.h"

class Server
{
public:
  Server()
    : mShutdown(false)
    , mListenSocket(0)
    , mTp( CreateThreadpool(0) )
  {
    InitializeThreadpoolEnvironment(&mTpEnv);
    SetThreadpoolCallbackPool(&mTpEnv, mTp);

    WSADATA wsaData;
    WSAStartup( MAKEWORD(2,2), &wsaData );
  }

  ~Server()
  {
    Shutdown(); // In case this wasn't called

    if (mListenSocket)
      delete mListenSocket;

    WSACleanup();

    CloseThreadpool(mTp);
    DestroyThreadpoolEnvironment(&mTpEnv);
  }

  void Shutdown() 
  {
    mShutdown = true;

    if (mListenSocket)
      mListenSocket->Close();
  }

  static void CALLBACK ReceivePacket(PTP_CALLBACK_INSTANCE instance, PVOID param)
  {
    RUdpSocket* client = reinterpret_cast<RUdpSocket*>(param);

    client->OnPacketReceive();
    client->Release();
  }

  void Listen(char* address, int port) 
  {
    ClientList clients;

    mListenSocket = new RUdpSocket();
    mListenSocket->Bind(address, port);

    // Maximum size of a UDP packet. More likely to be anywhere between 1.5k and 8k
    // but when using localhost, Windows increases the packet size to I think 16k or 32k (at least with TCP)
    const int MAX_UDP_PACKET = 65 * 1024;
    char buffer[MAX_UDP_PACKET];

    std::cout << "Server listening on " << address << ":" << port << std::endl;

    while (!mShutdown)
    {
      int bytes = mListenSocket->Receive(buffer, MAX_UDP_PACKET);

      if (bytes != SOCKET_ERROR && bytes != INVALID_SOCKET && bytes > 0)
      {
        // Check to see if this client has already connected.
        RUdpSocket* clientSocket = clients.FindSession( mListenSocket->HashEndPoint() );

        // Create a new Session object
        if (!clientSocket)
        {
          clientSocket = new RUdpSocket(*mListenSocket, mListenSocket->GetRemoteEndpoint(), mListenSocket->GetLocalEndpoint());

          clients.AddSession(clientSocket);
          
          // Every UDP socket will have a timer. When the timer expires it checks the age of the connection
          // and if the timestamp hasn't been updated, it will remove itself from the connection list
          clients.CreateCleanupTimer(clientSocket, &mTpEnv);
        }

        RUdpPacket* packet = new RUdpPacket;
        packet->FromBuffer(buffer, bytes);

        clientSocket->Push(packet);

        if (TrySubmitThreadpoolCallback(ReceivePacket, clientSocket, &mTpEnv) == FALSE)
          _ASSERT(0);

        Sleep(1);
      }
    }

    clients.RemoveAllSessions();
  }

private:
  bool mShutdown;
  RUdpSocket* mListenSocket;
  TP_CALLBACK_ENVIRON mTpEnv;
  PTP_POOL mTp;
};
